Example Tracking Pipeline

Get the model that will produce new inferences.

from chariot.models import get_model_by_id

model_id = "<The id of the model inferring over imagery>"
m = get_model_by_id(model_id)

Define a function to process inferences into the format the object tracking service expects.

from uuid import uuid4
from chariot import tracker as tracking_service

def process_detections(detections: list[dict]) -> list[tracking_service.UnitlessBox]:
    boxes: list[tracking_service.UnitlessBox] = []
    for d in detections:
        for label, score, (ymin, xmin, ymax, xmax) in zip(d["detection_classes"], d["detection_scores"], d["detection_boxes"]):
            width = xmax - xmin
            height = ymax - ymin
            cx = xmin + width/2
            cy = ymin + height/2
            box = tracking_service.UnitlessBox(
                id=str(uuid4()),
                label=label,
                score=float(score),
                x=float(cx),
                y=float(cy),
                w=float(width),
                h=float(height),
                covariance=[[1,0,0,0],[0,1,0,0],[0,0,1,0],[0,0,0,1]]
            )
            boxes.append(box)
    return boxes

Create a function that creates a tracker, reads a video, gets inferences, and updates the tracker.

import cv2
from PIL import Image
from datetime import datetime, UTC

from chariot import tracker as tracking_service

project_id = "<Project to create tracker in>"

def session(video_path: str, session_id: str):
    t = tracking_service.create_tracker(tracking_service.NewCreateTrackerRequest(
        project_id=project_id,
        kind=tracking_service.TrackerKind.BOX_UNITLESS,
        max_missing_updates=1,
        min_lifetime_before_active=1,
        assignment_function=tracking_service.TrackerAssignmentFunction.INTERSECTION_OVER_UNION,
        assignment_threshold=0.9,
    ))
    tracker_id = t.tracker_id
    print(f"Session ID: {session_id} | Tracker ID: {tracker_id}")
    
    cap = cv2.VideoCapture(video_path)
    sqn = 0
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        img = Image.fromarray(frame_rgb)
        metadata = {"session_id": session_id, "sequence_number": sqn}
        inference = m.detect(img, custom_metadata=metadata, return_inference_id=True)
        # Ensure the inference id matches the one in the inference store
        inference_id = inference["inference_id"] + "-0"
        boxes = process_detections(inference["inference"])
        tracking_service.update_tracker(tracking_service.NewUpdateTrackerRequest(
            tracker_id=tracker_id,
            external_input=tracking_service.ExternalTrackerInput(
                project_id=project_id,
                model_id=model_id,
                inference_id=inference_id,
                ts=datetime.now(UTC).isoformat(),
                sequence_number=sqn,
            ),
            boxes_unitless=boxes,
        ))
        sqn += 1
    return tracker_id

Define a function to get the inferences and tracks produced for the tracker.

from chariot.inference_store import inference
from chariot.inference_store import models
from chariot.inference_store import track

def track_generator(tracker_id: str, session_id: str, page_size: int = 50):
    offset = 0
    while True:
        inferences = inference.filter_inferences(model_id=model_id, request_body=models.NewGetInferencesRequest(
            filters=models.BaseInferenceFilter(
                data_source_filter=data_source, metadata_filter= [
                    models.MetadataFilter(
                        key="session_id",
                        type=models.MetadataFilterType.STRING,
                        operator=models.MetadataFilterOperator.EQUAL,
                        value=session_id,
                    ) 
                ]
            ),
            presign=True,
            pagination=models.Pagination(
                sort=models.PaginationSortField.CREATED_AT,
                direction=models.PaginationSortDirection.ASCENDING,
                offset=offset,
                limit=page_size,
            )
        ))
        if len(inferences) == 0:
            break
        
        out = {}
        for inference in inferences:
            tracks = track.get_tracks_for_inference(tracker_id=tracker_id, inference_id=inference.inference_id)
            out[inference.inference_id] = {"inference": inference, "tracks": tracks}

        yield out
        offset += page_size

Define a run method to produce inferences and tracks and to later collect them from the Inference Store for further processing.

from uuid import uuid4

def run_session(video_path: str):
    session_id = str(uuid4())
    tracker_id = session(video_path=video_path, session_id=session_id)
    for kv in track_generator(tracker_id=tracker_id, session_id=session_id):
      pass